package com.example.demo.aiInterface;

import com.example.demo.dao.ChatMessageDao;
import com.example.demo.entity.pojo.ChatMessage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.util.retry.Retry;

import java.time.Duration;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 基础AI助手实现类，包含共用的流处理、消息格式化和数据库操作逻辑
 * 特定的助手实现类可以继承此类，只需实现特定的业务逻辑
 */
public abstract class BaseAssistantImpl {

    protected final Logger logger = LoggerFactory.getLogger(this.getClass());
    protected static final int MAX_RETRIES = 3;
    protected static final Duration RETRY_DELAY = Duration.ofSeconds(2);

    @Autowired
    protected ChatMessageDao chatMessageDao;

    /**
     * 处理用户消息并获取AI响应
     * 子类需要实现此方法来提供特定类型的AI处理
     * 
     * @param sessionId 会话ID
     * @param userMessage 用户消息
     * @return AI响应的Flux流
     */
    protected abstract Flux<String> getAiResponse(String sessionId, String userMessage);
    
    /**
     * 获取助手类型名称，用于日志记录
     * 
     * @return 助手类型名称
     */
    protected abstract String getAssistantTypeName();

    /**
     * 处理聊天请求的主方法，包含通用的流处理和数据库操作
     * 
     * @param sessionId 会话ID
     * @param userMessage 用户消息
     * @return 格式化并控制速率后的AI响应流
     */
    public Flux<String> chat(String sessionId, String userMessage) {
        String assistantType = getAssistantTypeName();
        
        // 用于收集完整响应的容器
        AtomicReference<StringBuilder> responseCollector = new AtomicReference<>(new StringBuilder());

        logger.info("调用{}助手: sessionId={}, messagePreview={}",
                assistantType, sessionId, userMessage.length() > 30 ? userMessage.substring(0, 30) + "..." : userMessage);

        // 获取特定AI助手的响应流
        Flux<String> aiResponse = getAiResponse(sessionId, userMessage)
                .doOnSubscribe(s -> logger.info("开始请求{}AI回复", assistantType))
                .doOnComplete(() -> logger.info("{}AI回复完成", assistantType))
                .doOnError(error -> logger.error("{}AI回复出错: {}", assistantType, error.getMessage()))
                // 将流序列进行处理，优化输出格式
                .map(this::formatResponse)
                .retryWhen(Retry.backoff(MAX_RETRIES, RETRY_DELAY)
                        .doBeforeRetry(signal ->
                                logger.info("重试{}AI请求 (尝试 {}/{})",
                                        assistantType, signal.totalRetries() + 1, MAX_RETRIES))
                        .onRetryExhaustedThrow((spec, signal) -> signal.failure()))
                .onErrorResume(e -> {
                    // 当所有重试都失败时，返回一个友好的错误消息
                    logger.error("{}AI请求失败，已达到最大重试次数", assistantType, e);
                    return Flux.just("很抱歉，我现在遇到了一些问题。请稍后再试或者换一种提问方式。");
                });

        // 转换流，添加速率控制并保存完整响应
        return aiResponse
                .transform(this::applyStreamRateControl)
                .doOnNext(fragment -> {
                    // 将片段添加到收集器
                    responseCollector.get().append(fragment);
                })
                .doOnComplete(() -> {
                    // 流完成时，保存整个响应到数据库
                    String fullResponse = cleanupFullResponse(responseCollector.get().toString());
                    saveResponseToDatabase(sessionId, fullResponse);
                });
    }

    /**
     * 简单的流式速率控制
     * 减缓输出速度，使消息显示更自然
     */
    protected Flux<String> applyStreamRateControl(Flux<String> original) {
        return original.concatMap(fragment -> {
            // 估算此片段中的令牌数量（粗略估计为字符数的1/4，最小为1）
            int estimatedTokens = Math.max(1, fragment.length() / 4);

            // 计算延迟时间 (ms)
            long delayMillis = Math.max(10, estimatedTokens * 20);

            // 对于代码块，添加额外的延迟以确保完整性
            if (fragment.contains("```")) {
                delayMillis = Math.max(delayMillis, 150); // 确保代码块至少有150ms延迟
            }

            // 对于段落结束，添加自然的停顿
            if (fragment.endsWith("\n\n") || fragment.endsWith("。") ||
                fragment.endsWith("？") || fragment.endsWith("！")) {
                delayMillis = Math.max(delayMillis, 100); // 确保段落结束有100ms延迟
            }

            // 应用延迟并返回原始片段
            return Mono.just(fragment).delayElement(Duration.ofMillis(delayMillis));
        });
    }

    /**
     * 异步保存响应到数据库
     */
    protected void saveResponseToDatabase(String sessionId, String fullResponse) {
        // 使用Schedulers.boundedElastic()在后台线程执行
        Mono.fromCallable(() -> {
            try {
                logger.info("保存完整{}响应到数据库: sessionId={}, responseLength={}",
                        getAssistantTypeName(), sessionId, fullResponse.length());

                ChatMessage chatMessage = new ChatMessage();
                chatMessage.setSessionId(sessionId);
                chatMessage.setMessageType(1); // 1表示AI消息
                chatMessage.setMessageContent(fullResponse);
                chatMessage.setSendTime(System.currentTimeMillis());
                chatMessage.setStatus(1); // 已发送状态

                chatMessageDao.insert(chatMessage);
                logger.info("{}回复已成功保存到数据库", getAssistantTypeName());
                return true;
            } catch (Exception e) {
                logger.error("保存{}回复到数据库失败", getAssistantTypeName(), e);
                return false;
            }
        })
        .subscribeOn(Schedulers.boundedElastic())
        .subscribe();
    }

    /**
     * 格式化单个响应片段
     */
    protected String formatResponse(String response) {
        if (response == null || response.isEmpty()) {
            return "";
        }

        // 移除多余的标点符号
        response = response.replaceAll("([，,。.；;！!？?])\\1+", "$1");

        // 优化代码块格式
        response = formatCodeBlocks(response);

        return response;
    }

    /**
     * 清理完整的响应文本
     */
    protected String cleanupFullResponse(String response) {
        if (response == null || response.isEmpty()) {
            return "";
        }

        // 标准化换行符
        response = response.replaceAll("\\r\\n", "\n");

        // 标准化代码块前后的空行
        response = response.replaceAll("```([a-zA-Z]*)\\s*\n", "```$1\n");
        response = response.replaceAll("\n\\s*```", "\n```");

        // 移除重复的空行
        response = response.replaceAll("\\n{3,}", "\n\n");

        // 移除不必要的句末空格
        response = response.replaceAll("\\s+\n", "\n");

        return response;
    }

    /**
     * 格式化代码块
     */
    protected String formatCodeBlocks(String text) {
        // 识别未封闭的代码块并尝试修复
        if (text.contains("```") && !isCodeBlockClosed(text)) {
            return text; // 在流中不修改未封闭的代码块，等待后续内容
        }

        return text;
    }

    /**
     * 检查代码块是否正确闭合
     */
    protected boolean isCodeBlockClosed(String text) {
        Pattern pattern = Pattern.compile("```");
        Matcher matcher = pattern.matcher(text);
        int count = 0;
        while (matcher.find()) {
            count++;
        }
        return count % 2 == 0; // 如果是偶数，说明代码块已闭合
    }
} 